package top.zhangjianyong.tools.indicator.strategy.impl;

import org.springframework.stereotype.Component;
import top.zhangjianyong.tools.entity.ETFData;
import top.zhangjianyong.tools.indicator.context.IndicatorContext;
import top.zhangjianyong.tools.indicator.model.TechnicalIndicatorResult;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.List;

/**
 * RSI（相对强弱指数）指标计算器
 * 使用Wilder's RSI算法，窗口14
 *
 * @author zhangjianyong
 * @date 2024/01/01
 */
@Component
public class RSIIndicator extends AbstractIndicator {

    private static final String INDICATOR_NAME = "RSI";
    private static final int PERIOD = 14;

    @Override
    public String getName() {
        return INDICATOR_NAME;
    }

    @Override
    public int getMinDataSize() {
        return PERIOD;
    }

    @Override
    protected TechnicalIndicatorResult doCalculate(IndicatorContext context) {
        List<ETFData> dataList = context.getDataList();
        BigDecimal rsi = calculateWilderRSI(dataList);

        TechnicalIndicatorResult result = TechnicalIndicatorResult.success(getName());
        if (rsi != null) {
            result.addValue("rsi", rsi);
            result.addValue("rsi_14", rsi);
        } else {
            return TechnicalIndicatorResult.failure(getName(), "RSI计算失败，数据不足或数据异常");
        }

        return result;
    }

    /**
     * 使用Wilder's RSI算法计算相对强弱指数
     * Wilder's RSI使用指数平滑而不是简单平均，让最近的数据有更大的权重
     *
     * @param dataList ETF历史数据列表，按时间顺序排列（早到晚）
     * @return RSI值，范围0-100，null表示计算失败
     */
    private BigDecimal calculateWilderRSI(List<ETFData> dataList) {
        if (dataList.size() < PERIOD) {
            return null;
        }

        // 计算前PERIOD天的初始平均上涨和下跌
        BigDecimal initialAvgGain = BigDecimal.ZERO;
        BigDecimal initialAvgLoss = BigDecimal.ZERO;

        for (int i = 0; i < PERIOD; i++) {
            ETFData data = dataList.get(i);
            BigDecimal dailyRate = data.getDailyRate();

            if (dailyRate.compareTo(BigDecimal.ZERO) > 0) {
                initialAvgGain = initialAvgGain.add(dailyRate);
            } else if (dailyRate.compareTo(BigDecimal.ZERO) < 0) {
                initialAvgLoss = initialAvgLoss.add(dailyRate.abs());
            }
        }

        // 计算初始平均值
        initialAvgGain = initialAvgGain.divide(new BigDecimal(PERIOD), 8, RoundingMode.HALF_UP);
        initialAvgLoss = initialAvgLoss.divide(new BigDecimal(PERIOD), 8, RoundingMode.HALF_UP);

        // 如果初始平均下跌为0，说明前PERIOD天都是上涨，RSI为100
        if (initialAvgLoss.compareTo(BigDecimal.ZERO) == 0) {
            return new BigDecimal(100);
        }

        // 使用Wilder's平滑公式更新平均值（从第PERIOD+1天开始）
        BigDecimal currentAvgGain = initialAvgGain;
        BigDecimal currentAvgLoss = initialAvgLoss;

        // 平滑系数：13/14 (N-1)/N，其中N=PERIOD
        BigDecimal smoothFactor = new BigDecimal(PERIOD - 1);
        BigDecimal period = new BigDecimal(PERIOD);

        for (int i = PERIOD; i < dataList.size(); i++) {
            ETFData data = dataList.get(i);
            BigDecimal dailyRate = data.getDailyRate();

            BigDecimal todayGain = BigDecimal.ZERO;
            BigDecimal todayLoss = BigDecimal.ZERO;

            if (dailyRate.compareTo(BigDecimal.ZERO) > 0) {
                todayGain = dailyRate;
            } else if (dailyRate.compareTo(BigDecimal.ZERO) < 0) {
                todayLoss = dailyRate.abs();
            }

            // Wilder's平滑公式：
            // 新平均值 = (前一日平均值 × (N-1) + 今日值) / N
            currentAvgGain = currentAvgGain.multiply(smoothFactor)
                .add(todayGain)
                .divide(period, 8, RoundingMode.HALF_UP);

            currentAvgLoss = currentAvgLoss.multiply(smoothFactor)
                .add(todayLoss)
                .divide(period, 8, RoundingMode.HALF_UP);
        }

        // 如果当前平均下跌为0，RSI为100
        if (currentAvgLoss.compareTo(BigDecimal.ZERO) == 0) {
            return new BigDecimal(100);
        }

        // 计算RS值
        BigDecimal rs = currentAvgGain.divide(currentAvgLoss, 8, RoundingMode.HALF_UP);

        // 计算RSI值：RSI = 100 - (100 / (1 + RS))
        BigDecimal rsi = new BigDecimal(100)
            .subtract(new BigDecimal(100)
                .divide(rs.add(BigDecimal.ONE), 8, RoundingMode.HALF_UP));

        return rsi.setScale(2, RoundingMode.HALF_UP);
    }
}

